#![no_std]
#![doc = "PWM driver for the Rockchip RK3588 SoC channels used on OrangePi 5 Plus."]

use core::ptr::{read_volatile, write_volatile};

pub mod regs;

use regs::{ctrl, offset};

/// Errors returned by the PWM driver.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Error {
    /// The requested period evaluates to zero ticks.
    ZeroPeriod,
    /// Duty cycle cannot exceed the configured period.
    DutyExceedsPeriod,
    /// The nanosecond value cannot be represented with 32-bit ticks.
    TickOverflow,
    /// Oneshot mode accepts counts in the range 1..=256.
    InvalidOneshotCount,
}

/// PWM output polarity.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Polarity {
    /// Active (duty) portion drives high, inactive drives low.
    Normal,
    /// Active (duty) portion drives low, inactive drives high.
    Inverted,
}

/// PWM alignment selection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Alignment {
    /// Left-aligned waveform (edge aligned).
    Left,
    /// Center-aligned waveform.
    Center,
}

/// PWM run mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RunMode {
    /// Continuous operation.
    Continuous,
    /// Oneshot operation, repeating the pattern for the given number of
    /// periods.
    Oneshot(u16),
}

impl RunMode {
    /// Extract the oneshot repetition count, if any.
    #[inline]
    const fn oneshot_count(self) -> Option<u16> {
        match self {
            RunMode::Oneshot(count) => Some(count),
            RunMode::Continuous => None,
        }
    }
}

/// High-level configuration expressed in nanoseconds.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Config {
    /// PWM period in nanoseconds.
    pub period_ns: u64,
    /// PWM duty cycle in nanoseconds.
    pub duty_ns: u64,
    /// Output polarity mode.
    pub polarity: Polarity,
    /// Alignment mode.
    pub alignment: Alignment,
    /// Run mode.
    pub mode: RunMode,
    /// Whether the PWM should be enabled after configuration.
    pub enable: bool,
}

impl Config {
    /// Convenience constructor with sensible defaults.
    pub const fn new(period_ns: u64, duty_ns: u64) -> Self {
        Self {
            period_ns,
            duty_ns,
            polarity: Polarity::Normal,
            alignment: Alignment::Left,
            mode: RunMode::Continuous,
            enable: true,
        }
    }
}

/// Configuration expressed directly in timer ticks.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TickConfig {
    /// Period in controller ticks.
    pub period: u32,
    /// Duty cycle in controller ticks.
    pub duty: u32,
    /// Output polarity mode.
    pub polarity: Polarity,
    /// Alignment mode.
    pub alignment: Alignment,
    /// Run mode.
    pub mode: RunMode,
    /// Whether the PWM should be enabled after configuration.
    pub enable: bool,
}

impl TickConfig {
    /// Convenience constructor with defaults.
    pub const fn new(period: u32, duty: u32) -> Self {
        Self {
            period,
            duty,
            polarity: Polarity::Normal,
            alignment: Alignment::Left,
            mode: RunMode::Continuous,
            enable: true,
        }
    }
}

/// Snapshot of the controller state.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct State {
    /// Current counter value.
    pub counter: u32,
    /// Programmed period in ticks.
    pub period: u32,
    /// Programmed duty cycle in ticks.
    pub duty: u32,
    /// Whether output is currently enabled.
    pub enabled: bool,
    /// Effective polarity.
    pub polarity: Polarity,
    /// Effective alignment.
    pub alignment: Alignment,
    /// Effective run mode.
    pub mode: RunMode,
}

impl State {
    /// Converts the stored period to nanoseconds given controller properties.
    pub fn period_ns(&self, clock_hz: u64, prescaler: u32) -> u64 {
        ticks_to_ns(self.period as u64, clock_hz, prescaler)
    }

    /// Converts the stored duty to nanoseconds given controller properties.
    pub fn duty_ns(&self, clock_hz: u64, prescaler: u32) -> u64 {
        ticks_to_ns(self.duty as u64, clock_hz, prescaler)
    }
}

/// Rockchip PWM channel abstraction.
pub struct RockchipPwm {
    base_addr: usize,
    clock_hz: u64,
    prescaler: u32,
}

impl RockchipPwm {
    /// Creates a new PWM wrapper with the RK3588 default prescaler (1).
    pub const fn new(base_addr: usize, clock_hz: u64) -> Self {
        Self {
            base_addr,
            clock_hz,
            prescaler: regs::RK3588_PRESCALER,
        }
    }

    /// Creates a new PWM wrapper with an explicit prescaler.
    pub const fn with_prescaler(base_addr: usize, clock_hz: u64, prescaler: u32) -> Self {
        Self {
            base_addr,
            clock_hz,
            prescaler,
        }
    }

    /// Returns the controller clock in Hertz.
    #[inline]
    pub const fn clock_hz(&self) -> u64 {
        self.clock_hz
    }

    /// Returns the programmed prescaler.
    #[inline]
    pub const fn prescaler(&self) -> u32 {
        self.prescaler
    }

    /// Applies a configuration expressed in nanoseconds.
    pub fn configure(&self, cfg: &Config) -> Result<(), Error> {
        let period_ticks = self.ns_to_ticks(cfg.period_ns)?;
        let duty_ticks = self.ns_to_ticks(cfg.duty_ns)?;

        let tick_cfg = TickConfig {
            period: period_ticks,
            duty: duty_ticks,
            polarity: cfg.polarity,
            alignment: cfg.alignment,
            mode: cfg.mode,
            enable: cfg.enable,
        };

        self.configure_ticks(&tick_cfg)
    }

    /// Applies a configuration expressed in controller ticks.
    pub fn configure_ticks(&self, cfg: &TickConfig) -> Result<(), Error> {
        if cfg.period == 0 {
            return Err(Error::ZeroPeriod);
        }
        if cfg.duty > cfg.period {
            return Err(Error::DutyExceedsPeriod);
        }

        if let Some(count) = cfg.mode.oneshot_count() {
            if count == 0 || count > ctrl::ONESHOT_COUNT_MAX {
                return Err(Error::InvalidOneshotCount);
            }
        }

        let mut ctrl_val = self.read_reg(offset::CTRL);
        self.write_reg(offset::CTRL, ctrl_val | ctrl::LOCK_EN);
        ctrl_val |= ctrl::LOCK_EN;

        self.write_reg(offset::PERIOD, cfg.period);
        self.write_reg(offset::DUTY, cfg.duty);

        ctrl_val &= !ctrl::POLARITY_MASK;
        match cfg.polarity {
            Polarity::Normal => {
                ctrl_val |= ctrl::DUTY_POSITIVE;
                ctrl_val &= !ctrl::INACTIVE_POSITIVE;
            }
            Polarity::Inverted => {
                ctrl_val &= !ctrl::DUTY_POSITIVE;
                ctrl_val |= ctrl::INACTIVE_POSITIVE;
            }
        }

        match cfg.alignment {
            Alignment::Left => ctrl_val &= !ctrl::OUTPUT_CENTER,
            Alignment::Center => ctrl_val |= ctrl::OUTPUT_CENTER,
        }

        // Configure run mode.
        ctrl_val &= !ctrl::ONESHOT_COUNT_MASK;
        match cfg.mode {
            RunMode::Continuous => {
                ctrl_val |= ctrl::CONTINUOUS;
            }
            RunMode::Oneshot(count) => {
                debug_assert!(count >= 1);
                ctrl_val &= !ctrl::CONTINUOUS;
                ctrl_val |= ((count - 1) as u32) << ctrl::ONESHOT_COUNT_SHIFT;
            }
        }

        ctrl_val &= !ctrl::LOCK_EN;
        self.write_reg(offset::CTRL, ctrl_val);

        if cfg.enable {
            self.set_enabled(true);
        } else {
            self.set_enabled(false);
        }

        Ok(())
    }

    /// Enables or disables the PWM output.
    pub fn set_enabled(&self, enable: bool) {
        let mut ctrl_val = self.read_reg(offset::CTRL);

        if enable {
            let oneshot_active = (ctrl_val & ctrl::ONESHOT_COUNT_MASK) != 0;
            if oneshot_active {
                ctrl_val &= !ctrl::CONTINUOUS;
            } else {
                ctrl_val |= ctrl::CONTINUOUS;
            }
            ctrl_val |= ctrl::ENABLE;
        } else {
            ctrl_val &= !ctrl::ENABLE;
        }

        self.write_reg(offset::CTRL, ctrl_val);
    }

    /// Convenience wrapper to enable the PWM.
    #[inline]
    pub fn enable(&self) {
        self.set_enabled(true);
    }

    /// Convenience wrapper to disable the PWM.
    #[inline]
    pub fn disable(&self) {
        self.set_enabled(false);
    }

    /// Reads back the controller state.
    pub fn state(&self) -> State {
        let ctrl_val = self.read_reg(offset::CTRL);
        let period = self.read_reg(offset::PERIOD);
        let duty = self.read_reg(offset::DUTY);
        let counter = self.read_reg(offset::COUNTER);

        let polarity = if (ctrl_val & ctrl::DUTY_POSITIVE) != 0 {
            Polarity::Normal
        } else {
            Polarity::Inverted
        };

        let alignment = if (ctrl_val & ctrl::OUTPUT_CENTER) != 0 {
            Alignment::Center
        } else {
            Alignment::Left
        };

        let mode = if (ctrl_val & ctrl::ONESHOT_COUNT_MASK) != 0 {
            let raw = ((ctrl_val & ctrl::ONESHOT_COUNT_MASK) >> ctrl::ONESHOT_COUNT_SHIFT) as u16;
            RunMode::Oneshot(raw + 1)
        } else {
            RunMode::Continuous
        };

        State {
            counter,
            period,
            duty,
            enabled: (ctrl_val & ctrl::ENABLE) != 0,
            polarity,
            alignment,
            mode,
        }
    }

    /// Converts timer ticks back to nanoseconds.
    pub fn ticks_to_ns(&self, ticks: u64) -> u64 {
        ticks_to_ns(ticks, self.clock_hz, self.prescaler)
    }

    fn ns_to_ticks(&self, ns: u64) -> Result<u32, Error> {
        if ns == 0 {
            return Ok(0);
        }

        let prescaler = self.prescaler as u128;
        let clock = self.clock_hz as u128;
        let numerator = clock.checked_mul(ns as u128).ok_or(Error::TickOverflow)?;
        let denominator = prescaler * 1_000_000_000u128;

        let ticks = (numerator + denominator / 2) / denominator;
        if ticks == 0 {
            return Err(Error::ZeroPeriod);
        }
        if ticks > u32::MAX as u128 {
            return Err(Error::TickOverflow);
        }
        Ok(ticks as u32)
    }

    #[inline]
    fn read_reg(&self, offset: usize) -> u32 {
        unsafe { read_volatile((self.base_addr + offset) as *const u32) }
    }

    #[inline]
    fn write_reg(&self, offset: usize, value: u32) {
        unsafe { write_volatile((self.base_addr + offset) as *mut u32, value) }
    }
}

fn ticks_to_ns(ticks: u64, clock_hz: u64, prescaler: u32) -> u64 {
    if ticks == 0 {
        return 0;
    }
    let ticks = ticks as u128;
    let prescaler = prescaler as u128;
    let clock = clock_hz as u128;

    let numerator = ticks * prescaler * 1_000_000_000u128;
    ((numerator + clock / 2) / clock) as u64
}
